#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
from dotenv import load_dotenv
load_dotenv()
import os
import jwt
import datetime
from typing import Union
from rest_framework.exceptions import APIException
from apps.utils.exceptions import InvalidTokenException, TokenExpiredException


class TokenHelper:
    JWT_SECRET_KEY = os.environ.get('JWT_SECRET_KEY')
    JWT_ALGORITHM = 'HS256'
    JWT_EXPIRATION_DELTA = datetime.timedelta(days=int(os.environ.get("JWT_EXPIRATION_DELTA", 7)))

    @staticmethod
    def generate_jwt_token(user_id: str, wx_open_id: str):
        payload = {
            'wx_open_id': wx_open_id,
            'user_id': user_id,
            'exp': datetime.datetime.utcnow() + TokenHelper.JWT_EXPIRATION_DELTA
        }
        return jwt.encode(payload, TokenHelper.JWT_SECRET_KEY, algorithm=TokenHelper.JWT_ALGORITHM)

    @staticmethod
    def verify_token(token) -> Union[dict, None, APIException]:
        try:
            payload = jwt.decode(token, TokenHelper.JWT_SECRET_KEY, algorithms=[TokenHelper.JWT_ALGORITHM])
            return payload
        except jwt.ExpiredSignatureError as e:
            print('jwt.ExpiredSignatureError', e)
            raise TokenExpiredException()
        except jwt.InvalidTokenError as e:
            print('jwt.InvalidTokenError', e)
            print(e)
            raise InvalidTokenException()


if __name__ == '__main__':
    print(TokenHelper.generate_jwt_token('1', '2'))